import cv2
import numpy as np


def shift_mask(mask, ZY):
    if ZY == "ZY20210800001":
        dx, dy = 0, 0
    elif (ZY == "ZY20210800002") or (ZY == "ZY20210800004"):
        dx, dy = 15, -3
    elif ZY == "ZY20210800003":
        dx, dy = 15, -30
    else:
        raise NotImplementedError

    rows, cols, _ = mask.shape
    MAT = np.float32([[1, 0, dx], [0, 1, dy]])
    dst = cv2.warpAffine(mask, MAT, (cols, rows))
    return dst


def get_color_map(N=256):
    def bitget(byteval, idx):
        return ((byteval & (1 << idx)) != 0)

    cmap = np.zeros((N, 3), dtype=np.uint8)
    for i in range(N):
        r = g = b = 0
        c = i
        for j in range(8):
            r = r | (bitget(c, 0) << 7-j)
            g = g | (bitget(c, 1) << 7-j)
            b = b | (bitget(c, 2) << 7-j)
            c = c >> 3

        cmap[i] = np.array([r, g, b])

    return cmap
